import asyncio
import os
import dotenv

from openai import AsyncOpenAI

from agents import (
    Agent,
    Model,
    ModelProvider,
    OpenAIChatCompletionsModel,
    RunConfig,
    Runner,
    InputGuardrail,
    GuardrailFunctionOutput,
    set_tracing_disabled,
)
from pydantic import BaseModel

dotenv.load_dotenv()

client = AsyncOpenAI(api_key=os.getenv("AI_API_KEY"),base_url=os.getenv("AI_BASE_URL"))
set_tracing_disabled(disabled=True)

class CustomModelProvider(ModelProvider):
    def get_model(self, model_name: str | None) -> Model:
        return OpenAIChatCompletionsModel(model=os.getenv("AI_API_MODEL"), openai_client=client)

CUSTOM_MODEL_PROVIDER = CustomModelProvider()

# 分类结果模型
class ClassificationResult(BaseModel):
    question_type: str  # "history" 或 "math"
    reasoning: str

# 分类Agent
classification_agent = Agent(
    name="问题分类器",
    instructions="""
你是一个问题分类器。请分析用户问题并判断类型：

如果问题涉及历史（朝代、皇帝、历史人物、历史事件），返回：
{"question_type": "history", "reasoning": "这是历史问题"}

如果问题涉及数学（计算、方程、几何、公式），返回：
{"question_type": "math", "reasoning": "这是数学问题"}

只输出JSON格式，不要其他内容。
""",
    output_type=ClassificationResult,
)

async def classification_guardrail(ctx, agent, input_data):
    print(f"正在分析问题: {input_data}")
    result = await Runner.run(
        starting_agent=classification_agent, 
        input=input_data, 
        context=ctx.context,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER)
    )
    classification = result.final_output_as(ClassificationResult)
    print(f"分类结果: {classification.question_type} - {classification.reasoning}")
    
    return GuardrailFunctionOutput(
        output_info=classification,
        tripwire_triggered=False,
    )

# 定义专门的辅导员Agents
历史辅导员 = Agent(
    name="历史辅导员",
    instructions="""
    你是一个历史辅导员，专门回答历史问题。
    每一步都要解释你的推理过程并包含例子。
    在回答前加上"历史辅导员："
    """,
    handoff_description="专门处理历史问题的 agent"
)

数学辅导员 = Agent(
    name="数学辅导员", 
    instructions="""
    你是一个数学辅导员，专门回答数学问题。
    每一步都要解释你的推理过程并包含例子。
    在回答前加上"数学辅导员："
    """,
    handoff_description="专门处理数学问题的 agent"
)

# 智能分类Agent - 使用上下文路由信息
triage_agent = Agent(
    name="智能分类Agent",
    instructions="""
我是智能问题路由器。

当我接收到问题时，我会检查上下文中的分类信息：

如果上下文显示 classification="history"，我会立即将问题转交给历史辅导员。
如果上下文显示 classification="math"，我会立即将问题转交给数学辅导员。

我绝不会自己回答问题，只负责准确的路由转交。

根据分类结果，我现在需要转交给相应的专业辅导员。
""",
    handoffs=[历史辅导员, 数学辅导员],
    input_guardrails=[
        InputGuardrail(guardrail_function=classification_guardrail),
    ],
)

async def main():
    user_input = input("请你要问的问题: ")
    
    # 第一步：分类
    print(f"正在分析问题: {user_input}")
    classification_result = await Runner.run(
        starting_agent=classification_agent,
        input=user_input,
        run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
    )
    
    classification = classification_result.final_output_as(ClassificationResult)
    print(f"分类结果: {classification.question_type} - {classification.reasoning}")
    
    # 第二步：根据分类结果直接路由
    if classification.question_type == "history":
        print(">>> 转交给历史辅导员处理")
        result = await Runner.run(
            starting_agent=历史辅导员,
            input=user_input,
            run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
        )
    elif classification.question_type == "math":
        print(">>> 转交给数学辅导员处理")
        result = await Runner.run(
            starting_agent=数学辅导员,
            input=user_input,
            run_config=RunConfig(model_provider=CUSTOM_MODEL_PROVIDER),
        )
    else:
        print(">>> 未识别的问题类型")
        return
    
    print(result.final_output)

if __name__ == "__main__":  
    asyncio.run(main())


